{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Info Extraction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "it's much more easier to extract information of model from pytorch module than onnx...onnx doesn't have output shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "graph torch-jit-export (\n",
      "  %1[FLOAT, 1x3x224x224]\n",
      ") initializers (\n",
      "  %2[FLOAT, 64x3x3x3]\n",
      "  %3[FLOAT, 64]\n",
      "  %4[FLOAT, 64x64x3x3]\n",
      "  %5[FLOAT, 64]\n",
      "  %6[FLOAT, 128x64x3x3]\n",
      "  %7[FLOAT, 128]\n",
      "  %8[FLOAT, 128x128x3x3]\n",
      "  %9[FLOAT, 128]\n",
      "  %10[FLOAT, 256x128x3x3]\n",
      "  %11[FLOAT, 256]\n",
      "  %12[FLOAT, 256x256x3x3]\n",
      "  %13[FLOAT, 256]\n",
      "  %14[FLOAT, 256x256x3x3]\n",
      "  %15[FLOAT, 256]\n",
      "  %16[FLOAT, 256x256x3x3]\n",
      "  %17[FLOAT, 256]\n",
      "  %18[FLOAT, 512x256x3x3]\n",
      "  %19[FLOAT, 512]\n",
      "  %20[FLOAT, 512x512x3x3]\n",
      "  %21[FLOAT, 512]\n",
      "  %22[FLOAT, 512x512x3x3]\n",
      "  %23[FLOAT, 512]\n",
      "  %24[FLOAT, 512x512x3x3]\n",
      "  %25[FLOAT, 512]\n",
      "  %26[FLOAT, 512x512x3x3]\n",
      "  %27[FLOAT, 512]\n",
      "  %28[FLOAT, 512x512x3x3]\n",
      "  %29[FLOAT, 512]\n",
      "  %30[FLOAT, 512x512x3x3]\n",
      "  %31[FLOAT, 512]\n",
      "  %32[FLOAT, 512x512x3x3]\n",
      "  %33[FLOAT, 512]\n",
      "  %34[FLOAT, 4096x25088]\n",
      "  %35[FLOAT, 4096]\n",
      "  %36[FLOAT, 4096x4096]\n",
      "  %37[FLOAT, 4096]\n",
      "  %38[FLOAT, 1000x4096]\n",
      "  %39[FLOAT, 1000]\n",
      ") {\n",
      "  %41 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%1, %2)\n",
      "  %42 = Add[axis = 1, broadcast = 1](%41, %3)\n",
      "  %43 = Relu(%42)\n",
      "  %45 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%43, %4)\n",
      "  %46 = Add[axis = 1, broadcast = 1](%45, %5)\n",
      "  %47 = Relu(%46)\n",
      "  %48 = MaxPool[kernel_shape = [2, 2], pads = [0, 0], strides = [2, 2]](%47)\n",
      "  %50 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%48, %6)\n",
      "  %51 = Add[axis = 1, broadcast = 1](%50, %7)\n",
      "  %52 = Relu(%51)\n",
      "  %54 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%52, %8)\n",
      "  %55 = Add[axis = 1, broadcast = 1](%54, %9)\n",
      "  %56 = Relu(%55)\n",
      "  %57 = MaxPool[kernel_shape = [2, 2], pads = [0, 0], strides = [2, 2]](%56)\n",
      "  %59 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%57, %10)\n",
      "  %60 = Add[axis = 1, broadcast = 1](%59, %11)\n",
      "  %61 = Relu(%60)\n",
      "  %63 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%61, %12)\n",
      "  %64 = Add[axis = 1, broadcast = 1](%63, %13)\n",
      "  %65 = Relu(%64)\n",
      "  %67 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%65, %14)\n",
      "  %68 = Add[axis = 1, broadcast = 1](%67, %15)\n",
      "  %69 = Relu(%68)\n",
      "  %71 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%69, %16)\n",
      "  %72 = Add[axis = 1, broadcast = 1](%71, %17)\n",
      "  %73 = Relu(%72)\n",
      "  %74 = MaxPool[kernel_shape = [2, 2], pads = [0, 0], strides = [2, 2]](%73)\n",
      "  %76 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%74, %18)\n",
      "  %77 = Add[axis = 1, broadcast = 1](%76, %19)\n",
      "  %78 = Relu(%77)\n",
      "  %80 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%78, %20)\n",
      "  %81 = Add[axis = 1, broadcast = 1](%80, %21)\n",
      "  %82 = Relu(%81)\n",
      "  %84 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%82, %22)\n",
      "  %85 = Add[axis = 1, broadcast = 1](%84, %23)\n",
      "  %86 = Relu(%85)\n",
      "  %88 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%86, %24)\n",
      "  %89 = Add[axis = 1, broadcast = 1](%88, %25)\n",
      "  %90 = Relu(%89)\n",
      "  %91 = MaxPool[kernel_shape = [2, 2], pads = [0, 0], strides = [2, 2]](%90)\n",
      "  %93 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%91, %26)\n",
      "  %94 = Add[axis = 1, broadcast = 1](%93, %27)\n",
      "  %95 = Relu(%94)\n",
      "  %97 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%95, %28)\n",
      "  %98 = Add[axis = 1, broadcast = 1](%97, %29)\n",
      "  %99 = Relu(%98)\n",
      "  %101 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%99, %30)\n",
      "  %102 = Add[axis = 1, broadcast = 1](%101, %31)\n",
      "  %103 = Relu(%102)\n",
      "  %105 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%103, %32)\n",
      "  %106 = Add[axis = 1, broadcast = 1](%105, %33)\n",
      "  %107 = Relu(%106)\n",
      "  %108 = MaxPool[kernel_shape = [2, 2], pads = [0, 0], strides = [2, 2]](%107)\n",
      "  %109 = Reshape[shape = [1, -1]](%108)\n",
      "  %112 = Gemm[alpha = 1, beta = 1, broadcast = 1, transB = 1](%109, %34, %35)\n",
      "  %113 = Relu(%112)\n",
      "  %115, %116 = Dropout[is_test = 1, ratio = 0.5](%113)\n",
      "  %119 = Gemm[alpha = 1, beta = 1, broadcast = 1, transB = 1](%115, %36, %37)\n",
      "  %120 = Relu(%119)\n",
      "  %122, %123 = Dropout[is_test = 1, ratio = 0.5](%120)\n",
      "  %126 = Gemm[alpha = 1, beta = 1, broadcast = 1, transB = 1](%122, %38, %39)\n",
      "  return %126\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import onnx\n",
    "\n",
    "# Load the ONNX model\n",
    "model = onnx.load(\"onnx/vgg19.onnx\")\n",
    "\n",
    "# Check that the IR is well formed\n",
    "onnx.checker.check_model(model)\n",
    "\n",
    "# Print a human readable representation of the graph\n",
    "print(onnx.helper.printable_graph(model.graph))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python2.7/dist-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n"
     ]
    }
   ],
   "source": [
    "#import onnx_caffe2.backend as backend\n",
    "import onnx_tf.backend as backend\n",
    "import numpy as np\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract Shape of Initial Input Object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "rep = backend.prepare(model, device=\"CUDA:0\") # or \"CPU\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rep.input_dict['2']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_init_shape_dict(rep):\n",
    "    d = {}\n",
    "    for key in rep.input_dict:\n",
    "        tensor = rep.input_dict[key]\n",
    "        shape = np.array(tensor.shape, dtype=int)\n",
    "        d.update({key:shape})\n",
    "    return d\n",
    "shape_dict = get_init_shape_dict(rep)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(shape_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "pickle.dump(shape_dict,open('onnx/vgg19_init_shape_dict.pkl','wb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract Shape of Outputed Object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "init_shape_dict = pickle.load(open('onnx/vgg19_init_shape_dict.pkl','rb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "39"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(init_shape_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_output_shape_of_node(node, shape_dict, device = \"CPU\"):# or \"CUDA:0\"\n",
    "    \n",
    "    out_idx = node.output[0]\n",
    "    input_list = node.input # e.g. ['1', '2']\n",
    "    \n",
    "    inps = []\n",
    "    for inp_idx in input_list:\n",
    "        inp_shape = shape_dict[inp_idx] \n",
    "        rand_inp = np.random.random(size=inp_shape).astype('float16')\n",
    "        inps.append(rand_inp)\n",
    "    try:\n",
    "        out = backend.run_node(node=node, inputs=inps, device=device)\n",
    "        out_shape = out[0].shape \n",
    "    except:\n",
    "        out_shape = shape_dict[input_list[0]]\n",
    "        print(\"Op: [{}] run_node error! return inp_shape as out_shape\".format(node.op_type))\n",
    "        \n",
    "    return out_shape, out_idx "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test \n",
    "get_output_shape_of_node(model.graph.node[0], init_shape_dict.copy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test Op: [Dropout]\n",
    "shape_dict1 = init_shape_dict.copy()\n",
    "shape_dict1.update({'120':[1,3,4,5]})\n",
    "get_output_shape_of_node(model.graph.node[59], shape_dict1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test Op: [Gemm]\n",
    "shape_dict1 = init_shape_dict.copy()\n",
    "shape_dict1.update({'122':[1,4096]})\n",
    "for i, node in enumerate(model.graph.node[60:]):\n",
    "    st=time.time()\n",
    "    out_shape, out_idx = get_output_shape_of_node(node, shape_dict1)\n",
    "    shape_dict1.update({out_idx:out_shape})\n",
    "    print(\"out_shape: {} for Obj[{}], node [{}][{}]...{:.2f} sec\".format(out_shape, out_idx, i, node.op_type,time.time()-st))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# takes time!!!\n",
    "# 雖然可以自己算，怕網路複雜時會出錯，故還是用onnx backend去跑出output_shape(會多花點時間)\n",
    "def get_overall_shape_dict(init_shape_dict):\n",
    "    shape_dict = init_shape_dict.copy()\n",
    "    for i, node in enumerate(model.graph.node):\n",
    "        st=time.time()\n",
    "        out_shape, out_idx = get_output_shape_of_node(node, shape_dict)\n",
    "        shape_dict.update({out_idx:out_shape})\n",
    "        print(\"out_shape: {} for Obj[{}], node [{}][{}]...{:.2f} sec\".format(out_shape, out_idx, i, node.op_type,time.time()-st))\n",
    "    return shape_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python2.7/dist-packages/onnx_tf/backend.py:677: UserWarning: Unsupported kernel_shape attribute by Tensorflow in Conv operator. The attribute will be ignored.\n",
      "  UserWarning)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "out_shape: (1, 64, 224, 224) for Obj[41], node [0][Conv]...0.43 sec\n",
      "out_shape: (1, 64, 224, 224) for Obj[42], node [1][Add]...7.51 sec\n",
      "out_shape: (1, 64, 224, 224) for Obj[43], node [2][Relu]...7.49 sec\n",
      "out_shape: (1, 64, 224, 224) for Obj[45], node [3][Conv]...8.04 sec\n",
      "out_shape: (1, 64, 224, 224) for Obj[46], node [4][Add]...8.25 sec\n",
      "out_shape: (1, 64, 224, 224) for Obj[47], node [5][Relu]...8.45 sec\n",
      "out_shape: (1, 64, 112, 112) for Obj[48], node [6][MaxPool]...9.18 sec\n",
      "out_shape: (1, 128, 112, 112) for Obj[50], node [7][Conv]...4.09 sec\n",
      "out_shape: (1, 128, 112, 112) for Obj[51], node [8][Add]...5.81 sec\n",
      "out_shape: (1, 128, 112, 112) for Obj[52], node [9][Relu]...6.03 sec\n",
      "out_shape: (1, 128, 112, 112) for Obj[54], node [10][Conv]...6.48 sec\n",
      "out_shape: (1, 128, 112, 112) for Obj[55], node [11][Add]...6.49 sec\n",
      "out_shape: (1, 128, 112, 112) for Obj[56], node [12][Relu]...6.56 sec\n",
      "out_shape: (1, 128, 56, 56) for Obj[57], node [13][MaxPool]...6.60 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[59], node [14][Conv]...4.92 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[60], node [15][Add]...5.12 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[61], node [16][Relu]...5.31 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[63], node [17][Conv]...6.59 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[64], node [18][Add]...5.47 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[65], node [19][Relu]...5.42 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[67], node [20][Conv]...7.05 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[68], node [21][Add]...5.82 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[69], node [22][Relu]...5.82 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[71], node [23][Conv]...7.53 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[72], node [24][Add]...6.06 sec\n",
      "out_shape: (1, 256, 56, 56) for Obj[73], node [25][Relu]...6.16 sec\n",
      "out_shape: (1, 256, 28, 28) for Obj[74], node [26][MaxPool]...6.29 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[76], node [27][Conv]...7.69 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[77], node [28][Add]...5.49 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[78], node [29][Relu]...5.81 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[80], node [30][Conv]...10.89 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[81], node [31][Add]...5.81 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[82], node [32][Relu]...5.89 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[84], node [33][Conv]...11.13 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[85], node [34][Add]...6.18 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[86], node [35][Relu]...6.28 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[88], node [36][Conv]...11.81 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[89], node [37][Add]...6.68 sec\n",
      "out_shape: (1, 512, 28, 28) for Obj[90], node [38][Relu]...6.75 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[91], node [39][MaxPool]...6.83 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[93], node [40][Conv]...11.39 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[94], node [41][Add]...6.26 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[95], node [42][Relu]...6.26 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[97], node [43][Conv]...11.78 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[98], node [44][Add]...6.59 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[99], node [45][Relu]...6.58 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[101], node [46][Conv]...12.03 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[102], node [47][Add]...6.96 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[103], node [48][Relu]...6.87 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[105], node [49][Conv]...12.33 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[106], node [50][Add]...7.22 sec\n",
      "out_shape: (1, 512, 14, 14) for Obj[107], node [51][Relu]...7.23 sec\n",
      "out_shape: (1, 512, 7, 7) for Obj[108], node [52][MaxPool]...7.24 sec\n",
      "out_shape: (1, 25088) for Obj[109], node [53][Reshape]...7.16 sec\n",
      "out_shape: (1, 4096) for Obj[112], node [54][Gemm]...249.89 sec\n",
      "out_shape: (1, 4096) for Obj[113], node [55][Relu]...18.73 sec\n",
      "Op: [Dropout] run_node error! return inp_shape as out_shape\n",
      "out_shape: (1, 4096) for Obj[115], node [56][Dropout]...18.59 sec\n",
      "out_shape: (1, 4096) for Obj[119], node [57][Gemm]...57.49 sec\n",
      "out_shape: (1, 4096) for Obj[120], node [58][Relu]...20.28 sec\n",
      "Op: [Dropout] run_node error! return inp_shape as out_shape\n",
      "out_shape: (1, 4096) for Obj[122], node [59][Dropout]...20.19 sec\n",
      "out_shape: (1, 1000) for Obj[126], node [60][Gemm]...29.77 sec\n"
     ]
    }
   ],
   "source": [
    "overall_shape_dict = get_overall_shape_dict(init_shape_dict)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(overall_shape_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "pickle.dump(overall_shape_dict,open('onnx/vgg19_overall_shape_dict.pkl','wb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get Input/Output/Kernel Shape for Conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "overall_shape_dict = pickle.load(open('onnx/vgg19_overall_shape_dict.pkl','rb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_kernel_shape_dict():\n",
    "    conv_d = {}\n",
    "    for i, node in enumerate(model.graph.node):\n",
    "        if node.op_type == 'Conv':\n",
    "            for attr in node.attribute:\n",
    "                if attr.name == \"kernel_shape\":\n",
    "                    kernel_shape = np.array(attr.ints, dtype=int)\n",
    "                    break\n",
    "            inp_idx = node.input[0]\n",
    "            out_idx = node.output[0]\n",
    "            inp_shape = overall_shape_dict[inp_idx]\n",
    "            out_shape = overall_shape_dict[out_idx]\n",
    "            conv_d.update({i:(inp_idx, out_idx, inp_shape, out_shape, kernel_shape)})\n",
    "            print(\"for node [{}][{}]:\\ninp_shape: {} from obj[{}], \\nout_shape: {} from obj[{}], \\nkernel_shape: {} \\n\"\n",
    "                  .format(i, node.op_type, inp_shape, inp_idx, out_shape, out_idx, kernel_shape ))\n",
    "    return conv_d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "for node [0][Conv]:\n",
      "inp_shape: [  1   3 224 224] from obj[1], \n",
      "out_shape: (1, 64, 224, 224) from obj[41], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [3][Conv]:\n",
      "inp_shape: (1, 64, 224, 224) from obj[43], \n",
      "out_shape: (1, 64, 224, 224) from obj[45], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [7][Conv]:\n",
      "inp_shape: (1, 64, 112, 112) from obj[48], \n",
      "out_shape: (1, 128, 112, 112) from obj[50], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [10][Conv]:\n",
      "inp_shape: (1, 128, 112, 112) from obj[52], \n",
      "out_shape: (1, 128, 112, 112) from obj[54], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [14][Conv]:\n",
      "inp_shape: (1, 128, 56, 56) from obj[57], \n",
      "out_shape: (1, 256, 56, 56) from obj[59], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [17][Conv]:\n",
      "inp_shape: (1, 256, 56, 56) from obj[61], \n",
      "out_shape: (1, 256, 56, 56) from obj[63], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [20][Conv]:\n",
      "inp_shape: (1, 256, 56, 56) from obj[65], \n",
      "out_shape: (1, 256, 56, 56) from obj[67], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [23][Conv]:\n",
      "inp_shape: (1, 256, 56, 56) from obj[69], \n",
      "out_shape: (1, 256, 56, 56) from obj[71], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [27][Conv]:\n",
      "inp_shape: (1, 256, 28, 28) from obj[74], \n",
      "out_shape: (1, 512, 28, 28) from obj[76], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [30][Conv]:\n",
      "inp_shape: (1, 512, 28, 28) from obj[78], \n",
      "out_shape: (1, 512, 28, 28) from obj[80], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [33][Conv]:\n",
      "inp_shape: (1, 512, 28, 28) from obj[82], \n",
      "out_shape: (1, 512, 28, 28) from obj[84], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [36][Conv]:\n",
      "inp_shape: (1, 512, 28, 28) from obj[86], \n",
      "out_shape: (1, 512, 28, 28) from obj[88], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [40][Conv]:\n",
      "inp_shape: (1, 512, 14, 14) from obj[91], \n",
      "out_shape: (1, 512, 14, 14) from obj[93], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [43][Conv]:\n",
      "inp_shape: (1, 512, 14, 14) from obj[95], \n",
      "out_shape: (1, 512, 14, 14) from obj[97], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [46][Conv]:\n",
      "inp_shape: (1, 512, 14, 14) from obj[99], \n",
      "out_shape: (1, 512, 14, 14) from obj[101], \n",
      "kernel_shape: [3 3] \n",
      "\n",
      "for node [49][Conv]:\n",
      "inp_shape: (1, 512, 14, 14) from obj[103], \n",
      "out_shape: (1, 512, 14, 14) from obj[105], \n",
      "kernel_shape: [3 3] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "conv_d = get_kernel_shape_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate n_param, n_flops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_d = pickle.load(open('onnx/vgg19_conv_shape_dict.pkl','rb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(20024384, 19523280896)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n_param = 0\n",
    "n_flops = 0\n",
    "for k in conv_d:\n",
    "    #i:(inp_idx, out_idx, inp_shape, out_shape, kernel_shape)\n",
    "    inp_shape, out_shape, kernel_shape = conv_d[k][2],conv_d[k][3],conv_d[k][4]\n",
    "    h,w,c,n,H,W = kernel_shape[1], kernel_shape[1], inp_shape[1], out_shape[1], out_shape[2], out_shape[3]\n",
    "    n_param  += n*(h*w*c+1)\n",
    "    n_flops  += H*W*n*(h*w*c+1)\n",
    "n_param, n_flops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
